import re
import json
from typing import Dict, Tuple
from volcengine.visual.VisualService import VisualService
from exception import AppException
from common.enums.apply import DrawsEnum
from common.utils.tools import ToolsUtil
from common.utils.config import ConfigUtil
from common.cache.pool_cache import KeyPoolCache


class DrawsMjServer:
    """ MJ绘画 """
    def __init__(self, params: Dict[str, any]) -> None:
        # 绘画渠道
        self.channel = "zsy"

        # 基础参数
        self.prompt: str = params.get("prompt", "")               # 提示词
        self.action: str = str(params.get("action", ""))          # 操作类型: [generate, upsample, variation]
        self.model: str = str(params.get("model", ""))            # 生成模式: [mj, niji]
        self.scale: str = str(params.get("scale", ""))            # 图片比例: [1:1, 1:2, ...]
        self.version: str = str(params.get("version", ""))        # 版本参数: [6.0, 5.2, ...]
        self.image_id: str = str(params.get("image_id", ""))      # 图片的ID
        self.image_ref: str = str(params.get("image_ref", ""))    # 参考图片
        self.callback: str = params.get("callback_url", "")       # 回调地址

        # 其它参数
        others = params.get("other_params") or {}
        self.mode: str = str(others.get("mode", ""))              # 生成模式: [fast, relax, turbo]
        self.style: str = str(others.get("style", ""))            # 风格样式: [default=动漫, cute=可爱, ...]
        self.no_content: str = str(others.get("no_content", ""))  # 排除内容: [去除图片的某个东西]
        self.other_params = {
            "mode": self.mode,
            "style": self.style,
            "no_content": self.no_content
        }

        # 高级参数
        self.complex_params = {
            # 种子参数
            "seed": params.get("complex_params", {}).get("seed", ""),
            # 图片权重: 默认1,[0.5-2],数值越高越接近参考图,反之偏向描述词
            "iw": float(params.get("complex_params", {}).get("iw", "1")),
            # 图片质量: 默认1,[0.25,1],图片细节刻画,越小出图越快
            "q": float(params.get("complex_params", {}).get("q", "1")),
            # 图片模糊: 默认100,[0-100],参数越高,风格化细节越多,生成时间越久
            "s": int(params.get("complex_params", {}).get("s", "100")),
            # 图片随机: 默认0,[0-100],数值越大随机性越高,较低则还是按照描述词走
            "c": int(params.get("complex_params", {}).get("c", "0"))
        }

        # 验证参数
        self._check_params()

        # 构建指令
        self._make_instruct()

    async def imagine(self):
        """ 生成图片 """
        config = await ConfigUtil.get("draws", "mj") or {}
        self.channel = config.get("channel", "")

        if self.channel == "zsy":
            return await self.post_zsy()
        elif self.channel == "acedata":
            return await self.post_aceData()
        else:
            raise AppException("无可用绘画通道,请管理员检查配置信息。")

    async def post_zsy(self):
        """
        知识云MJ

        Document:
            https://data.zhishuyun.com/documents

        Author:
            zero
        """
        url: str = "https://api.zhishuyun.com/midjourney/imagine"
        if self.mode == "relax":
            url += "/relax"
        elif self.mode == "turbo":
            url += "/turbo"

        token: str = await self.get_apikey()
        response = await ToolsUtil.post_curl(
            url=f"{url}?token=" + token,
            headers={
                "Accept": "application/json",
                "Content-Type": "application/json"
            },
            data={
                "action": self.action,
                "prompt": self.prompt,
                "image_id": self.image_id,
                "callback_url": self.callback,
                "translation": True,
                "timeout": 100
            })
        return self._handle_response(response)

    async def post_aceData(self):
        """
        AceData

        Document:

        """
        token: str = await self.get_apikey()
        response = await ToolsUtil.post_curl(
            url="https://api.acedata.cloud/midjourney/imagine",
            headers={
                "Accept": "application/json",
                "Content-Type": "application/json",
                "Authorization": "Bearer " + token
            },
            data={
                "mode": self.mode,
                "action": self.action,
                "prompt": self.prompt,
                "image_id": self.image_id,
                "mask": self.image_ref,
                "callback_url": self.callback,
                "translation": True,
                "timeout": 100
            })
        return self._handle_response(response)

    async def get_apikey(self) -> str:
        """ 获取密钥 """
        if self.channel == "acedata":
            marks = f"mj-{self.channel}"
        else:
            marks = f"mj-{self.channel}-{self.mode}"
        apikey = await KeyPoolCache(scene="draws", channel=marks).get_key() or {}
        if not apikey or not apikey.get("api_key"):
            raise AppException("尚未配置密钥,请与管理员联系")
        return apikey["api_key"]

    def _make_instruct(self):
        """ 构建指令 """
        # 参考图片
        if self.image_ref:
            self.prompt = self.image_ref + " " + self.prompt
        # 图片比例
        if self.scale:
            self.prompt += f" --ar {self.scale}"
        # 排除内容
        if self.no_content:
            self.prompt += f" --no {self.no_content}"
        # 版本参数
        if self.version:
            if self.version.startswith("--v") or self.version.startswith("--niji"):
                self.prompt += f" {self.version}"
                if self.style and self.style != "default":
                    self.prompt += f" --style {self.style}"
        # 高级参数
        if self.complex_params:
            complex_params = ""
            for k, v in self.complex_params.items():
                if not v:
                    continue
                complex_params += f" --{k} {v}"
            if complex_params:
                self.prompt += complex_params

    def _check_params(self):
        """ 验证参数 """

        if self.action not in ["generate", "upscale1", "upscale2", "upscale3", "upscale4",
                               "rerzoll", "variation1", "variation2", "variation3", "variation4"]:
            raise AppException("操作类型错误了")

        if self.action != "generate" and not self.image_id:
            raise AppException("请指定操作图片")

        if self.mode not in ["fast", "relax", "turbo"]:
            raise AppException("生成模式选择异常")

        if self.model not in ["mj", "niji"]:
            raise AppException("生成模型选择异常")

        if self.scale not in ["1:1", "1:2", "4:3", "9:16", "3:4", "16:9", "3:2", "2:3"]:
            raise AppException("不支持的图片比例")

        mj_version = (DrawsEnum.mj_config("mj_version") or {}).get("mj")
        if not mj_version and self.model == "mj" and self.version not in mj_version:
            raise AppException("模型版本选择异常")

        if self.model == "niji":
            mj_style = DrawsEnum.mj_config("mj_style") or {}
            niji_version = (DrawsEnum.mj_config("mj_version") or {}).get("niji")
            if self.version not in niji_version:
                raise AppException("模型版本选择异常")
            if self.style not in mj_style.keys():
                raise AppException("生成风格选择异常")

    @classmethod
    def _handle_response(cls, response):
        if response.get("code") == "invalid_token":
            raise Exception("令牌无效, 请确保您的密钥正确。")
        elif response.get("code") == "token_mismatched":
            raise Exception("令牌不匹配, 请确保您的密钥正确。")
        elif response.get("code") == "used_up":
            raise Exception("API套餐已经用完，请管理员到平台充值。")
        elif response.get("code") == "unknown":
            raise Exception("未知错误，需要联系API平台系统管理员。")
        elif response.get("code") == "unverified":
            raise Exception("未完成实名认证, 需管理员到API平台进行实名认证。")
        elif response.get("code") == "forbidden":
            raise Exception(response.get("detail", "forbidden"))
        return response


class DrawsDoubaoServer:
    """ 豆包绘画 """
    def __init__(self, params: Dict[str, any] = None) -> None:
        self.visual_service = None
        self.params = params
        self.complex_params = {}

    async def serv(self):
        apikey = await self.get_apikey()
        self.visual_service = VisualService()
        self.visual_service.set_ak(apikey.get("api_key"))
        self.visual_service.set_sk(apikey.get("api_secret"))
        return self.visual_service

    @classmethod
    async def get_apikey(cls) -> dict:
        """ 获取密钥 """
        apikey = await KeyPoolCache(scene="draws", channel="doubao").get_key() or {}
        if not apikey or not apikey.get("api_key"):
            raise AppException("尚未配置密钥,请与管理员联系")
        return apikey

    async def imagine(self):
        if self.params.get("model") == "high_aes_general_v21_L":
            return await self.text_to_image()
        elif self.params.get("model") == "high_aes":
            return await self.anime_to_image()
        elif self.params.get("model") == "i2i_xl_sft":
            return await self.image_to_image()
        else:
            raise AppException(f"绘画模型异常: {self.params.get('model')}")

    async def text_to_image(self):
        """
        文生图 (通用2.1-文生图)

        Document:
            https://www.volcengine.com/docs/6791/1366783

        Author:
            zero
        """
        await self.serv()

        try:
            width, height = self._scale_to_wh(self.params.get("scale", "512x512"))
            complex_params = self.params.get("complex_params", {})
            self.complex_params = {
                "seed": int(complex_params.get("seed") or -1),
                "scale": float(complex_params.get("scale") or 3.5),
                "ddim_steps": int(complex_params.get("ddim_steps") or 25)
            }

            resp = self.visual_service.cv_sync2async_submit_task({
                # 算法名称
                "req_key": "high_aes_general_v21_L",
                # 模型版本
                "model_version": "general_v2.1_L",
                # 美感配置
                "req_schedule_conf": "general_v20_9B_pe",
                # 提示词值
                "prompt": str(self.params.get("prompt")),
                # 随机种子: 默认值: [-1]
                "seed": self.complex_params["seed"],
                # 影响文本: 默认3.5, [1~10]
                "scale": self.complex_params["scale"],
                # 生成步数: 默认25, [1~200]
                "ddim_steps": self.complex_params["ddim_steps"],
                # 图像宽度
                "width": width,
                # 图像高度
                "height": height,
                # 文本扩写: [针对prompt优化]
                "use_pre_llm": True,
                # 超分功能: [True=文生图+AIGC超分, False=文生图]
                "use_sr": True
            })
            return self._handle_response(resp)
        except Exception as e:
            self._handle_exception(e)

    async def image_to_image(self):
        """
        图生图 (通用XL pro-图生图)

        Document:
            https://www.volcengine.com/docs/6791/1330203

        Author:
            zero
        """
        await self.serv()

        try:
            complex_params = self.params.get("complex_params", {})
            self.complex_params = {
                "seed": int(complex_params.get("seed") or -1),
                "scale": float(complex_params.get("scale") or 7.0),
                "ddim_steps": int(complex_params.get("ddim_steps") or 20),
            }

            resp = self.visual_service.cv_sync2async_submit_task({
                # 算法名称
                "req_key": "i2i_xl_sft",
                # 提示词值
                "prompt": self.params.get("prompt"),
                # 随机种子: 默认-1
                "seed": self.complex_params["seed"],
                # 影响文本: 默认7.0, 范围[1~30]
                "scale": self.complex_params["scale"],
                # 生成步数: 默认20, 范围[1~50]
                "ddim_steps": self.complex_params["ddim_steps"],
                # 参考图片
                "image_urls": [self.params.get("image_ref")] if self.params.get("image_ref", None) else []
            })
            return self._handle_response(resp)
        except Exception as e:
            self._handle_exception(e)

    async def anime_to_image(self):
        """
        动漫图 (动漫1.3.X-文生图/图生图)

        Document:
            https://www.volcengine.com/docs/6791/1213131

        Author:
            zero
        """
        await self.serv()

        try:
            width, height = self._scale_to_wh(self.params.get("scale", "512x512"))
            complex_params = self.params.get("complex_params", {})
            self.complex_params = {
                "seed": int(complex_params.get("seed") or -1),
                "scale": float(complex_params.get("scale", 7.0)),
                "strength": float(complex_params.get("strength", 0.7)),
                "ddim_steps": int(complex_params.get("ddim_steps", 25))
            }

            resp = self.visual_service.cv_sync2async_submit_task({
                # 算法名称
                "req_key": "high_aes",
                # 模型版本
                "model_version": "anime_v1.3.1",
                # 提示词值
                "prompt": self.params.get("prompt"),
                # 随机种子: 默认-1
                "seed": self.complex_params["seed"],
                # 影响文本: 默认7, [1~30]
                "scale": self.complex_params["scale"],
                # 接近原图, 默认0.7, [0~1.0]
                "strength": self.complex_params["strength"],
                # 生成步数: 默认25, [1~200]
                "ddim_steps": self.complex_params["ddim_steps"],
                # 参考图片
                "image_urls": [self.params.get("image_ref")] if self.params.get("image_ref", None) else [],
                # 图像宽度
                "width": width,
                # 图像高度
                "height": height
            })
            return self._handle_response(resp)
        except Exception as e:
            self._handle_exception(e)

    async def fetch(self, req_key: str, task_id: str):
        """
        获取生成结果

        Args:
            req_key (str): key
            task_id (str): 任务ID

        Returns:
            dict

        Author:
            zero
        """
        try:
            await self.serv()
            body = self.visual_service.cv_sync2async_get_result({
                "req_key": req_key,
                "task_id": task_id,
                "req_json": json.dumps({
                    "return_url": True
                })
            })
            return self._handle_response(body)
        except Exception as e:
            raise AppException(str(e))

    @classmethod
    def _scale_to_wh(cls, scale: str) -> Tuple[int, int]:
        """
        比例转宽高

        Args:
            scale (str): 512x512

        Returns:
            tuple: (width, height)

        Author:
            zero
        """
        if not scale:
            return 512, 512
        scales = scale.split("x")
        if len(scales) >= 2:
            return int(scales[0]), int(scales[1])
        return 512, 512

    @classmethod
    def _handle_response(cls, response):
        """
        处理响应结果

        Args:
            response (dict): {"code": 10000, "data": {"task_id": "7468667968", "message": "Success"}}

        Returns:
            dict:
                {"task_id": '7468667968814088203'}
                {"image_urls": ["https://xxx.png"]}
                {'binary_data_base64': [], 'image_urls': None, 'resp_data': '', 'status': 'in_queue'}

        Author:
            zero
        """
        if not response or response.get("status") != 10000 or not response.get("data"):
            code: str = str(response.get("code"))
            message: str = str(response.get("message"))
            raise AppException(f"{code}:{message}")

        return response.get("data")

    @classmethod
    def _handle_exception(cls, e):
        try:
            except_string = str(e).strip("b").strip("'")
            except_string = except_string.replace('\\"', '"')
            cleaned_string = re.sub(r',?"time_elapsed":".*?"', '', except_string)
            cleaned_string = re.sub(r',\s*}', '}', cleaned_string)
            cleaned_string = re.sub(r'{\s*,', '{', cleaned_string)
            error = json.loads(cleaned_string)
            if error.get("message"):
                code = str(error.get("code"))
                message = str(error.get("message"))
            else:
                wrong = error["ResponseMetadata"]["Error"]
                code = wrong.get("Code")
                message = wrong.get("Message")
            if code == "InvalidAccessKey":
                raise Exception("令牌无效, 请确保您的密钥正确。")
            elif code == "FlowLimitExceeded":
                raise Exception("请求过于频繁, 超出了限速(QPS), 请稍后重试。")
        except Exception as e:
            raise Exception(f"1: {str(e)}")
        raise Exception(f"{code}: {message}")


class DrawsDalleServer:
    """ Dalle绘画 """

    def __init__(self, params: Dict[str, any] = None):
        self.params = params
        self.baseUrl = "https://api.openai.com/v1"
        other_params = self.params.get("other_params", {})
        self.other_params = {
            "quality": other_params.get("quality") or "standard",
            "style": other_params.get("style") or "vivid"
        }

    async def imagine(self):
        # 获取配置
        config = await ConfigUtil.get("draws", "dalle") or {}
        if config.get("agency_api"):
            self.baseUrl = config["agency_api"].strip("/").strip()

        # 获取密钥
        token: str = await self.get_apikey()

        # 发起请求
        try:
            return await ToolsUtil.post_curl(
                url=f"{self.baseUrl}/images/generations",
                headers={
                    "Content-Type": "application/json",
                    "Authorization": "Bearer " + token
                },
                timeout=300,
                data={
                    "n": 1,
                    "model": self.params.get("model") or "dall-e-3",
                    "prompt": self.params.get("prompt"),
                    "size": self.params.get("scale") or "1:1",
                    "quality": self.other_params["quality"],
                    "style": self.other_params["style"],
                    "response_format": "b64_json"
                })
        except Exception as e:
            raise AppException(str(e))

    @classmethod
    async def get_apikey(cls) -> str:
        """ 获取密钥 """
        apikey = await KeyPoolCache(scene="draws", channel="dalle").get_key() or {}
        if not apikey or not apikey.get("api_key"):
            raise AppException("尚未配置密钥,请与管理员联系")
        return apikey["api_key"]


class DrawsSdServer:
    """ SD绘画 """

    def __init__(self):
        # 接收参数
        # 获取密钥
        pass

    async def imagine(self):
        data = {
            "prompt": "",
            "negative_prompt": "",
            "seed": "",
            "sampler_name": "",
            "steps": "",
            "cfg_scale": "",
            "batch_size": "",
            "width": "",
            "height": ""
        }

        # 文生图
        url = "/sdapi/v1/txt2img"

        # 图生图
        url = "sdapi/v1/img2img"

        # 发起请求
        await ToolsUtil.post_curl(
            url=url,
            headers={
                "Accept": "application/json",
                "Content-Type": "application/json"
            },
            data=data
        )
